"""
模型管理路由
"""
import os
import uuid
import shutil
from pathlib import Path
from typing import Optional, List
from fastapi import APIRouter, Depends, UploadFile, File, Form, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from sqlalchemy.orm import Session

from database import get_db
from schemas import (
    ModelCreateRequest, ModelUpdateRequest, ModelInfo,
    ModelTransferRequest, FineTuneRequest, GenericResponse, DataResponse, ListResponse,
    ModelEvaluationResultsResponse, EvaluationTaskListResponse, EvaluationTaskItem,
    ModelEvaluateRequest, ModelEvaluateResponse,
    ScenarioTrendRequest, ScenarioTrendResponse
)
from services import ModelService
from exceptions import NotFoundException, ValidationException
from response_utils import success_response, error_response, list_response
from config import config
from pathlib import Path
from sqlalchemy.orm import Session
from database import get_db, EvaluationRecord

router = APIRouter(prefix="/api/v1", tags=["模型管理"])

# 模型服务实例
model_service = ModelService()

@router.post("/models/upload", response_model=DataResponse[dict])
async def upload_model_file(
    file: UploadFile = File(..., description="模型文件")
):
    """上传模型文件到临时目录"""
    try:
        # 验证文件类型
        allowed_extensions = {'.pt', '.pth', '.pkl', '.onnx', '.h5', '.pb', '.tflite'}
        file_extension = Path(file.filename).suffix.lower()
        
        if file_extension not in allowed_extensions:
            raise HTTPException(
                status_code=400, 
                detail=f"不支持的文件格式，支持的格式: {', '.join(allowed_extensions)}"
            )
        
        # 验证文件大小（限制为500MB）
        max_size = 500 * 1024 * 1024  # 500MB
        file_content = await file.read()
        if len(file_content) > max_size:
            raise HTTPException(status_code=400, detail="文件大小不能超过500MB")
        
        # 生成唯一的文件名
        temp_filename = f"{uuid.uuid4()}{file_extension}"
        temp_file_path = config.paths.temp_dir / temp_filename
        
        # 确保temp目录存在
        config.paths.temp_dir.mkdir(parents=True, exist_ok=True)
        
        # 保存文件到临时目录
        with open(temp_file_path, "wb") as buffer:
            buffer.write(file_content)
        
        # 计算文件大小
        file_size = len(file_content)
        size_mb = file_size / (1024 * 1024)
        
        return success_response(
            data={
                "temp_path": str(temp_file_path),
                "original_filename": file.filename,
                "temp_filename": temp_filename,
                "file_size": file_size,
                "size_mb": round(size_mb, 2)
            },
            message="文件上传成功"
        )
        
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"文件上传失败: {str(e)}")

@router.post("/models", response_model=DataResponse[ModelInfo])
async def create_model(
    model: ModelCreateRequest,
    db: Session = Depends(get_db)
):
    """创建模型"""
    try:
        # 将Pydantic模型转换为字典
        model_dict = model.model_dump()
        
        # 如果提供了temp路径，处理文件移动
        temp_path = model_dict.get("weight_file_path")
        temp_file_to_delete = None
        
        if temp_path:
            # 如果路径是temp目录的文件，需要处理文件移动
            if temp_path.startswith(str(config.paths.temp_dir)):
                # 获取原始文件扩展名
                temp_file = Path(temp_path)
                if not temp_file.exists():
                    raise ValidationException(f"临时文件不存在: {temp_path}")
                
                temp_file_to_delete = temp_file  # 记录需要删除的临时文件
            else:
                # 对于非temp目录的路径，验证文件是否存在
                file_path = Path(temp_path)
                if not file_path.exists():
                    raise ValidationException(f"指定的模型文件不存在: {temp_path}")
        
        result = model_service.create_model(db, model_dict, temp_file_to_delete)
        
        # 检查是否是跳过的情况
        if result.get("skipped", False):
            return success_response(
                data=result["existing_model"], 
                message=result["message"]
            )
        
        return success_response(data=result, message="模型创建成功")
        
    except ValidationException as e:
        # 如果创建失败，清理临时文件
        if temp_file_to_delete and temp_file_to_delete.exists():
            try:
                os.remove(temp_file_to_delete)
            except:
                pass
        raise HTTPException(status_code=400, detail=e.message)
    except Exception as e:
        # 如果创建失败，清理临时文件
        if temp_file_to_delete and temp_file_to_delete.exists():
            try:
                os.remove(temp_file_to_delete)
            except:
                pass
        raise HTTPException(status_code=500, detail=f"模型创建失败: {str(e)}")

@router.get("/models", response_model=ListResponse[ModelInfo])
async def list_models(
    page: int = 1,
    page_size: int = 20,
    model_type: Optional[str] = None,
    status: Optional[str] = None,
    db: Session = Depends(get_db)
):
    """获取模型列表"""
    models, total_count = model_service.list_models(
        db, page=page, page_size=page_size, model_type=model_type, status=status
    )
    # 创建分页信息
    from response_utils import create_pagination_info
    pagination = create_pagination_info(page, page_size, total_count)
    
    # 调用服务层的统计方法
    summary = model_service.get_models_summary(db, model_type, status)
    
    return list_response(
        items=models,
        total_count=total_count,
        pagination=pagination,
        summary=summary,
        message="查询成功"
    )

@router.get("/models/{model_id}", response_model=DataResponse[ModelInfo])
async def get_model(model_id: str, db: Session = Depends(get_db)):
    """获取模型详情"""
    try:
        model = model_service.get_model_by_id(db, model_id)
        return success_response(data=model, message="查询成功")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)

@router.put("/models/{model_id}", response_model=DataResponse[ModelInfo])
async def update_model(
    model_id: str,
    model_update: ModelUpdateRequest,
    db: Session = Depends(get_db)
):
    """更新模型信息"""
    try:
        # 将Pydantic模型转换为字典
        update_dict = model_update.model_dump(exclude_unset=True)
        updated_model = model_service.update_model(db, model_id, update_dict)
        return success_response(data=updated_model, message="模型更新成功")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)
    except ValidationException as e:
        raise HTTPException(status_code=400, detail=e.message)

@router.post("/models/{model_id}/publish", response_model=DataResponse[dict])
async def publish_model(model_id: str, db: Session = Depends(get_db)):
    """发布模型"""
    try:
        result = model_service.publish_model(db, model_id)
        return success_response(data=result, message="模型发布成功")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)

@router.get("/models/{model_id}/path", response_model=DataResponse[dict])
async def get_model_path(model_id: str, db: Session = Depends(get_db)):
    """获取模型路径"""
    try:
        model_path = model_service.get_model_path(db, model_id)
        return success_response(data={"path": str(model_path)}, message="查询成功")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)

@router.post("/models/transfer", response_model=DataResponse[dict])
async def transfer_model(
    model_id: str = Form(...),
    target_node: str = Form(...),
    db: Session = Depends(get_db)
):
    """传输模型"""
    try:
        result = model_service.transfer_model(db, model_id, target_node)
        return success_response(data=result, message="模型传输成功")
    except (NotFoundException, ValidationException) as e:
        raise HTTPException(status_code=400, detail=e.message)

@router.post("/models/finetune", response_model=DataResponse[dict])
async def finetune_model(
    model_file: UploadFile = File(...),
    dataset_name: str = Form(...),
    epochs: int = Form(100),
    batch_size: int = Form(16),
    learning_rate: float = Form(0.001),
    db: Session = Depends(get_db)
):
    """微调模型"""
    try:
        result = model_service.finetune_model(
            db, model_file, dataset_name, epochs, batch_size, learning_rate
        )
        return success_response(data=result, message="模型微调启动成功")
    except (NotFoundException, ValidationException) as e:
        raise HTTPException(status_code=400, detail=e.message)

@router.get("/models/{model_id}/download")
async def download_model(model_id: str, db: Session = Depends(get_db)):
    """下载模型"""
    try:
        model_path = model_service.get_model_path(db, model_id)
        return FileResponse(
            path=str(model_path),
            filename=f"{model_id}.pt",
            media_type='application/octet-stream'
        )
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)

@router.get("/models/{model_id}/evaluation-results", response_model=DataResponse[ModelEvaluationResultsResponse])
async def get_model_evaluation_results(model_id: str, db: Session = Depends(get_db)):
    """
    获取模型评估结果
    
    通过模型ID找到对应的训练任务ID，并返回训练任务输出目录中的评估结果图片文件：
    - f1_curve.png
    - P_curve.png  
    - PR_curve.png
    - R_curve.png
    - results.png
    """
    try:
        result = model_service.get_model_evaluation_results(db, model_id)
        return success_response(data=result, message="获取模型评估结果成功")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)
    except ValidationException as e:
        raise HTTPException(status_code=400, detail=e.message)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取模型评估结果失败: {str(e)}")

@router.post("/models/{model_id}/evaluate", response_model=DataResponse[ModelEvaluateResponse])
async def evaluate_model(
    model_id: str,
    request: ModelEvaluateRequest,
    db: Session = Depends(get_db)
):
    """
    在线评测模型：
    - 指定数据集名称和样本数量等参数，对模型进行一次性评测
    - 返回评测指标与评测临时目录
    """
    try:
        result = model_service.evaluate_model(db, model_id, request)
        return success_response(data=result, message="模型评测完成")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)
    except ValidationException as e:
        raise HTTPException(status_code=400, detail=e.message)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型评测失败: {str(e)}")

@router.get("/evaluations/{eval_id}/{filename}")
async def get_evaluation_image(eval_id: str, filename: str):
    """访问评测临时目录中的可视化图片"""
    try:
        eval_root = config.paths.temp_dir / eval_id
        # 支持 preds、preds2 或任意以 preds 开头的目录
        candidate_dirs = [eval_root / "preds", eval_root / "preds2"]
        if eval_root.exists():
            for d in eval_root.iterdir():
                if d.is_dir() and d.name.startswith("preds"):
                    candidate_dirs.append(d)
        file_path = None
        for d in candidate_dirs:
            p = d / filename
            if p.exists():
                file_path = p
                break
        if file_path is None:
            raise HTTPException(status_code=404, detail="文件不存在")
        if not file_path.exists():
            raise HTTPException(status_code=404, detail="文件不存在")
        # 简单的安全检查，限制在 temp 目录下
        try:
            file_path.resolve().relative_to(config.paths.temp_dir.resolve())
        except ValueError:
            raise HTTPException(status_code=403, detail="访问被拒绝")
        # 推断媒体类型
        ext = file_path.suffix.lower()
        media_type = {
            ".png": "image/png",
            ".jpg": "image/jpeg",
            ".jpeg": "image/jpeg"
        }.get(ext, "application/octet-stream")
        return FileResponse(str(file_path), media_type=media_type, filename=filename)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"读取评测图片失败: {str(e)}")

@router.post("/models/{model_id}/evaluate/preds-dir", response_model=DataResponse[str])
async def evaluate_model_return_preds_dir(
    model_id: str,
    request: ModelEvaluateRequest,
    db: Session = Depends(get_db)
):
    """仅返回可视化输出目录(例如 preds2 的绝对路径)。"""
    try:
        result = model_service.evaluate_model(db, model_id, request)
        preds_dir = result.get("preds_dir")
        if not preds_dir:
            raise HTTPException(status_code=500, detail="未生成可视化输出目录")
        return success_response(data=preds_dir, message="ok")
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)
    except ValidationException as e:
        raise HTTPException(status_code=400, detail=e.message)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"模型评测失败: {str(e)}")

@router.get("/evaluation-records", response_model=ListResponse[dict])
async def list_evaluation_records(
    page: int = 1,
    page_size: int = 20,
    model_id: Optional[str] = None,
    dataset_name: Optional[str] = None,
    db: Session = Depends(get_db)
):
    """查询评测记录列表。"""
    try:
        query = db.query(EvaluationRecord)
        if model_id:
            query = query.filter(EvaluationRecord.model_id == model_id)
        if dataset_name:
            query = query.filter(EvaluationRecord.dataset_name == dataset_name)
        total = query.count()
        rows = query.order_by(EvaluationRecord.created_at.desc()) \
            .offset((page-1)*page_size).limit(page_size).all()
        items = []
        import json as _json
        for r in rows:
            items.append({
                "id": r.id,
                "eval_id": r.eval_id,
                "model_id": r.model_id,
                "model_name": r.model_name,
                "dataset_id": r.dataset_id,
                "dataset_name": r.dataset_name,
                "eval_dir": r.eval_dir,
                "preds_dir": r.preds_dir,
                "preds_subdir": r.preds_subdir,
                "used_samples": r.used_samples,
                "metrics": _json.loads(r.metrics) if r.metrics else {},
                "params": _json.loads(r.params) if r.params else {},
                "labels_in_eval": r.labels_in_eval,
                "max_label_class_id": r.max_label_class_id,
                "nc": r.nc,
                "device_type": r.device_type,
                "status": r.status,
                "created_at": r.created_at.strftime('%Y-%m-%d %H:%M:%S') if r.created_at else None
            })
        from response_utils import create_pagination_info
        pagination = create_pagination_info(page, page_size, total)
        return list_response(items=items, total_count=total, pagination=pagination, message="查询成功")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")

@router.get("/evaluation-records/{eval_id}", response_model=DataResponse[dict])
async def get_evaluation_record(eval_id: str, db: Session = Depends(get_db)):
    """按 eval_id 获取单条评测记录。"""
    try:
        rec = db.query(EvaluationRecord).filter(EvaluationRecord.eval_id == eval_id).first()
        if not rec:
            raise HTTPException(status_code=404, detail="评测记录不存在")
        import json as _json
        data = {
            "id": rec.id,
            "eval_id": rec.eval_id,
            "model_id": rec.model_id,
            "model_name": rec.model_name,
            "dataset_id": rec.dataset_id,
            "dataset_name": rec.dataset_name,
            "eval_dir": rec.eval_dir,
            "preds_dir": rec.preds_dir,
            "preds_subdir": rec.preds_subdir,
            "used_samples": rec.used_samples,
            "metrics": _json.loads(rec.metrics) if rec.metrics else {},
            "params": _json.loads(rec.params) if rec.params else {},
            "labels_in_eval": rec.labels_in_eval,
            "max_label_class_id": rec.max_label_class_id,
            "nc": rec.nc,
            "device_type": rec.device_type,
            "status": rec.status,
            "created_at": rec.created_at.strftime('%Y-%m-%d %H:%M:%S') if rec.created_at else None
        }
        return success_response(data=data, message="查询成功")
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")

@router.get("/evaluation-records/{eval_id}/urls", response_model=DataResponse[dict])
async def list_evaluation_image_urls(eval_id: str):
    """
    根据 eval_id 返回 preds2 目录下所有图片的可访问 URL 列表。
    若不存在 preds2，则回退到以 preds 开头的目录（如 preds）。
    """
    try:
        eval_root = config.paths.temp_dir / eval_id
        dir_candidates = []
        if eval_root.exists() and eval_root.is_dir():
            # 优先 preds2，其次任意以 preds 开头的目录
            if (eval_root / "preds2").exists():
                dir_candidates.append(eval_root / "preds2")
            for d in sorted(eval_root.iterdir()):
                if d.is_dir() and d.name.startswith("preds") and d not in dir_candidates:
                    dir_candidates.append(d)
        # 选择第一个可用目录，否则返回空列表但200
        target_dir = None
        for d in dir_candidates:
            if d.exists() and d.is_dir():
                target_dir = d
                break
        allowed_exts = {".png", ".jpg", ".jpeg"}
        urls = []
        chosen_dir = str(target_dir) if target_dir else None
        if target_dir:
            for p in sorted(target_dir.glob("*")):
                if p.is_file() and p.suffix.lower() in allowed_exts:
                    urls.append(f"/api/v1/evaluations/{eval_id}/{p.name}")
        # 返回候选目录，便于排查
        return success_response(
            data={
                "eval_id": eval_id,
                "urls": urls,
                "dir": chosen_dir,
                "dir_candidates": [str(d) for d in dir_candidates],
                "temp_root": str(config.paths.temp_dir),
            },
            message="查询成功"
        )
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")

@router.get("/evaluation-tasks", response_model=EvaluationTaskListResponse)
async def get_evaluation_task_list(
    page: int = 1,
    page_size: int = 20,
    db: Session = Depends(get_db)
):
    """
    获取评测任务列表
    
    仅返回有训练任务的模型的评测信息，包括：
    1. 任务名称（仅来自训练任务的模型有此字段）
    2. 模型名称
    3. 模型ID
    4. 数据集名称
    5. 模型精度（优先使用mAP@0.5，其次使用precision）
    6. 模型遗忘率（基于recall和precision计算：1 - F1_score）
    7. 模型的tag标签
    8. 模型的创建时间
    """
    try:
        # 获取评测任务列表
        evaluation_tasks, total_count = model_service.get_evaluation_task_list(
            db, page=page, page_size=page_size
        )
        
        # 创建分页信息
        from response_utils import create_pagination_info
        pagination = create_pagination_info(page, page_size, total_count)
        
        # 获取统计摘要
        summary = model_service.get_evaluation_task_summary(db)
        
        return EvaluationTaskListResponse(
            success=True,
            message="获取评测任务列表成功",
            data=evaluation_tasks,
            total_count=total_count,
            pagination=pagination,
            summary=summary
        )
        
    except ValidationException as e:
        raise HTTPException(status_code=400, detail=e.message)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取评测任务列表失败: {str(e)}")

@router.get("/models/{model_id}/evaluation-images/{filename}")
async def get_model_evaluation_image(model_id: str, filename: str, db: Session = Depends(get_db)):
    """获取模型评估结果图片文件"""
    try:
        # 支持的评估结果图片文件名
        allowed_files = {"F1_curve.png", "P_curve.png", "PR_curve.png", "R_curve.png", "results.png"}
        
        if filename not in allowed_files:
            raise HTTPException(
                status_code=400, 
                detail=f"不支持的文件名，支持的文件: {', '.join(allowed_files)}"
            )
        
        # 获取模型评估结果信息
        evaluation_result = model_service.get_model_evaluation_results(db, model_id)
        
        # 查找指定的文件
        target_file_info = None
        for file_info in evaluation_result["evaluation_files"]:
            if file_info["filename"] == filename:
                target_file_info = file_info
                break
        
        if not target_file_info:
            raise HTTPException(status_code=404, detail=f"评估结果文件不存在: {filename}")
        
        if not target_file_info["exists"]:
            raise HTTPException(status_code=404, detail=f"评估结果文件不存在: {filename}")
        
        # 构建文件路径
        file_path = Path(target_file_info["file_path"])
        
        # 检查文件是否存在
        if not file_path.exists():
            raise HTTPException(status_code=404, detail="评估结果图片文件不存在")
        
        # 安全检查：确保文件在模型目录内
        models_dir = config.paths.models_dir
        try:
            file_path.resolve().relative_to(models_dir.resolve())
        except ValueError:
            raise HTTPException(status_code=403, detail="访问被拒绝")
        
        # 根据文件扩展名确定媒体类型
        file_ext = file_path.suffix.lower()
        media_type_map = {
            '.png': 'image/png',
            '.jpg': 'image/jpeg',
            '.jpeg': 'image/jpeg',
            '.gif': 'image/gif',
            '.bmp': 'image/bmp',
            '.tiff': 'image/tiff',
            '.tif': 'image/tiff'
        }
        media_type = media_type_map.get(file_ext, 'image/png')
        
        # 返回图片文件
        return FileResponse(
            path=str(file_path),
            media_type=media_type,
            filename=filename
        )
        
    except HTTPException:
        raise
    except NotFoundException as e:
        raise HTTPException(status_code=404, detail=e.message)
    except ValidationException as e:
        raise HTTPException(status_code=400, detail=e.message)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取评估结果图片失败: {str(e)}")

@router.post("/models/scenario-trend", response_model=DataResponse[ScenarioTrendResponse])
async def get_scenario_trend(request: ScenarioTrendRequest):
    """
    获取场景趋势分析
    
    根据scenario字段返回遗忘率趋势和准确率趋势数据
    - scenario: 应用场景（人员安全、环境巡视、设备运检I、设备运检S）
    - 遗忘率趋势: 包含5个float值，呈下降趋势（从0.3到0.09）
    - 准确率趋势: 包含5个float值，呈上升趋势（从0.75到0.85）
    """
    try:
        # 根据scenario生成不同的趋势数据
        scenario_data = {
            "人员安全": {
                "forget": [0.30, 0.25, 0.18, 0.12, 0.09],
                "precision": [0.75, 0.78, 0.81, 0.83, 0.85]
            },
            "环境巡视": {
                "forget": [0.28, 0.22, 0.16, 0.11, 0.09],
                "precision": [0.76, 0.79, 0.82, 0.84, 0.85]
            },
            "设备运检I": {
                "forget": [0.30, 0.24, 0.19, 0.13, 0.09],
                "precision": [0.75, 0.77, 0.80, 0.82, 0.85]
            },
            "设备运检S": {
                "forget": [0.29, 0.23, 0.17, 0.12, 0.09],
                "precision": [0.75, 0.78, 0.81, 0.84, 0.85]
            }
        }
        
        # 获取对应场景的数据
        trend_data = scenario_data[request.scenario]
        
        # 创建响应对象
        response_data = ScenarioTrendResponse(
            scenario=request.scenario,
            forget=trend_data["forget"],
            precision=trend_data["precision"]
        )
        
        return success_response(data=response_data, message="场景趋势分析获取成功")
        
    except KeyError:
        raise HTTPException(status_code=400, detail=f"不支持的场景: {request.scenario}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取场景趋势分析失败: {str(e)}")

 